package test.demo;

import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;
import com.alibaba.fastjson.JSON;
import tech.waterism.util.JacksonUtil;

import java.util.Arrays;
import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 三水源滞后演算法（三水源汇流）
 * 分钟尺度
 * 系统模型集成
 *
 * @author 初京刚
 * @version 5.0
 */
public class LAG_3_min extends PredictBase {

    private double C0;
    private double C1;
    private double C2;
    private double CI;
    private double CG;
    private double X;
    private double KK;
    private double LAG;
    private double CS;
    private int MP;
    private double F;
    private Date[] TM;
    private double[][] R;
    private int T;
    private double[] QSIG;
    private double[] QXSIG;
    private double[] QSIG0;
    private double QSP;
    private double QIP;
    private double QGP;

    public LAG_3_min(Map<String, String> data) {
        super();

        CG = Double.parseDouble(data.get("CG"));
        CI = Double.parseDouble(data.get("CI"));
        X = Double.parseDouble(data.get("X"));
        KK = Double.parseDouble(data.get("KK"));
        LAG = Double.parseDouble(data.get("LAG"));
        CS = Double.parseDouble(data.get("CS"));
        MP = Integer.parseInt(data.get("MP"));
        F = Double.parseDouble(data.get("F"));

        TM = JSON.parseObject(data.get("dt"), Date[].class);
        R = JSON.parseObject(data.get("rSim"), double[][].class);
        T = Integer.parseInt(data.get("clen"));        // index.clen
        QSIG0 = JSON.parseObject(data.get("QSIG"), double[].class);
        QSP = Double.parseDouble(data.get("QSP"));
        QIP = Double.parseDouble(data.get("QIP"));
        QGP = Double.parseDouble(data.get("QGP"));
        QXSIG = new double[MP + 1];
        int I;
        for (I = 0; I < MP + 1; I++) {
            if (data.get("QXSIG") == null) {
                QXSIG[I] = 0.0;
            } else {
                if (I > JSON.parseObject(data.get("QXSIG"), double[].class).length - 1) {
                    QXSIG[I] = QXSIG[I - 1];
                } else {
                    QXSIG[I] = JSON.parseObject(data.get("QXSIG"), double[].class)[I];
                }
            }
        }
    }

    /**
     * 模型计算
     */
    public Map<String, Object> LAG3model() {
        double[] RS = R[0];
        double[] RI = R[1];
        double[] RG = R[2];
        int Pnum = RS.length;
        QSIG = new double[Pnum + (int) LAG];

        if (LAG <= 1) {
            QSIG[0] = QSIG0[0];
        } else {
            int I;
            for (I = 0; I < LAG; I++) {
                if (QSIG0.length >= LAG) {
                    QSIG[I] = QSIG0[I];
                } else {
                    if (I < QSIG0.length) {
                        QSIG[I] = QSIG0[I];
                    } else {
                        QSIG[I] = QSIG[I - 1];
                    }
                }
            }
        }

        double CP = F / T / 0.06;
        CG = Math.pow(CG, T / 1440.0);
        CI = Math.pow(CI, T / 1440.0);
        double FKT = KK - KK * X + 0.5 * T;
        C0 = (0.5 * T - KK * X) / FKT;
        C1 = (KK * X + 0.5 * T) / FKT;
        C2 = (KK - KK * X - 0.5 * T) / FKT;

        double QSIG1;
        double[][] QSIGS = new double[Pnum + 1][QSIG0.length];
        if (LAG <= 1) {
            QSIG1 = QSIG[0];
        } else {
            QSIG1 = QSIG[(int) (LAG - 1)];
        }

        for (int i = 0; i < QSIG0.length; i++) {
            QSIGS[0][i] = QSIG0[i];
        }

        double[][] QXSIGS = new double[Pnum + 1][MP + 1];
        double[] QXSIGSS = new double[MP + 1];
        for (int I = 0; I < MP + 1; I++) {
            QXSIGS[0][I] = QXSIG[I];
            QXSIGSS[I] = QXSIG[I];
        }

        double[] QG = new double[Pnum];
        double[] QS = new double[Pnum];
        double[] QI = new double[Pnum];

        double[] QGS = new double[Pnum + 1];
        double[] QSS = new double[Pnum + 1];
        double[] QIS = new double[Pnum + 1];
        QGS[0] = QGP;
        QSS[0] = QSP;
        QIS[0] = QIP;

        for (int J = 0; J < Pnum; J++) {
            QGP = QGP * CG + RG[J] * (1.0 - CG) * CP;
            QIP = QIP * CI + RI[J] * (1.0 - CI) * CP;
            QSP = RS[J] * CP;

            QG[J] = QGP;
            QS[J] = QSP;
            QI[J] = QIP;

            QGS[J + 1] = QGP;
            QSS[J + 1] = QSP;
            QIS[J + 1] = QIP;

            QSIG1 = QSIG1 * CS + (QGP + QIP + QSP) * (1.0 - CS);
            double QTSIG = QSIG1;
            QSIG[(int) (J + LAG)] = LCHCO(MP, QTSIG, QXSIGSS);
            for (int I = 0; I < MP + 1; I++) {
                QXSIGS[J + 1][I] = QXSIGSS[I];
            }

            Calendar time = Calendar.getInstance();
            time.setTime(TM[J]);
            Integer Mth = time.get(Calendar.MONTH) + 1; // 获得降雨开始月份
            Integer Yth = time.get(Calendar.YEAR);
            if (Yth > 2012 && Mth > 8) {
                QSIG[J] = QSIG[J] * 0.9;
            }
        }

        for (int i = 1; i <= QSIG.length - LAG; i++) {
            if (i < QSIG0.length) {
                for (int j = 0; j < QSIG0.length; j++) {
                    if (i + j < QSIG0.length) {
                        QSIGS[i][j] = QSIG0[j + i];
                    } else {
                        QSIGS[i][j] = QSIG[i + j - QSIG0.length];
                    }
                }
            } else {
                for (int j = 0; j < QSIG0.length; j++) {
                    if (i + j < QSIG.length) {
                        QSIGS[i][j] = QSIG[i + j];
                    } else {
                        QSIGS[i][j] = QSIG[QSIG.length - 1];
                    }
                }
            }
        }

        Map<String, Object> simulationResult = new HashMap<>();

        simulationResult.put("RS", R[0]);
        simulationResult.put("RI", R[1]);
        simulationResult.put("RG", R[2]);
        simulationResult.put("QS", QS);
        simulationResult.put("QI", QI);
        simulationResult.put("QG", QG);
        simulationResult.put("QSim", QSIG);
        simulationResult.put("QGS", QGS);
        simulationResult.put("QSS", QSS);
        simulationResult.put("QIS", QIS);
        simulationResult.put("QSIGS", QSIGS);
        simulationResult.put("QXSIGS", QXSIGS);

        return simulationResult;
    }

    private double LCHCO(double MP, double RQ, double[] QX) {
        double IM = MP + 1;
        if (IM == 1.0) {
            QX[(int) (IM - 1)] = RQ;
        } else {
            for (int J = 1; J < IM; J++) {
                double Q1 = RQ;
                double Q2 = QX[J - 1];
                double Q3 = QX[J];
                QX[J - 1] = RQ;
                RQ = C0 * Q1 + C1 * Q2 + C2 * Q3;
            }
            QX[(int) (IM - 1)] = RQ;
        }
        return RQ;
    }

    private List<String> toStringList(double[] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = String.valueOf(array0[i]);
        }
        return Arrays.asList(array);
    }

    private List<String> toStringList(double[][] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = Arrays.toString(array0[i]);
        }
        return Arrays.asList(array);
    }

    @Override
    public void init() {

    }

    @Override
    public PredictResult predict() {
        PredictResult result = new PredictResult();
        Map<String, Object> map = LAG3model();
        result.setQSim(JacksonUtil.convert(map.get("QSim"), double[].class));
        result.setRSim(new double[][]{JacksonUtil.convert(map.get("RS"), double[].class),
                JacksonUtil.convert(map.get("RI"), double[].class),
                JacksonUtil.convert(map.get("RG"), double[].class)});

        List<String> QXSIG = toStringList(JacksonUtil.convert(map.get("QXSIGS"), double[][].class));
        List<String> QSIG = toStringList(JacksonUtil.convert(map.get("QSIGS"), double[][].class));
        List<String> QSP = toStringList(JacksonUtil.convert(map.get("QSS"), double[].class));
        List<String> QIP = toStringList(JacksonUtil.convert(map.get("QIS"), double[].class));
        List<String> QGP = toStringList(JacksonUtil.convert(map.get("QGS"), double[].class));

        Map<String, List<String>> status = new HashMap<>();
        //河段下断面各时段初始流量
        status.put("QXSIG", QXSIG);
        //各时段初始总基流
        status.put("QSIG", QSIG);
        //地表坡面汇流
        status.put("QSP", QSP);
        //坡面汇流-壤中流
        status.put("QIP", QIP);
        //地下坡面汇流
        status.put("QGP", QGP);
        result.setNewStatus(status);

        return result;
    }
}